import os
import sys
import json
import requests
import tarfile
import base64
import io
from urllib.parse import urlparse
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor

# Docker Hub 的认证 URL
DOCKER_HUB_AUTH_URL = "https://auth.docker.io/token"
DOCKER_HUB_REGISTRY_URL = "registry-1.docker.io"

def get_auth_url(registry_url):
    if registry_url == DOCKER_HUB_REGISTRY_URL:
        return DOCKER_HUB_AUTH_URL
    if not registry_url.startswith("http://") and not registry_url.startswith("https://"):
        registry_url = f"https://{registry_url}"
    response = requests.get(f"{registry_url}/v2/")
    if response.status_code == 401:
        auth_header = response.headers.get('Www-Authenticate')
        if auth_header:
            auth_url = auth_header.split('realm="')[1].split('"')[0]
            return auth_url
    return None

def get_auth_token(auth_url, username, password, service, scope):
    if auth_url == DOCKER_HUB_AUTH_URL:
        params = {
            "service": service,
            "scope": scope
        }
        if username and password:
            auth = base64.b64encode(f"{username}:{password}".encode()).decode()
            headers = {
                "Authorization": f"Basic {auth}"
            }
            response = requests.get(auth_url, headers=headers, params=params)
        else:
            response = requests.get(auth_url, params=params)
    else:
        auth = base64.b64encode(f"{username}:{password}".encode()).decode()
        headers = {
            "Authorization": f"Basic {auth}"
        }
        params = {
            "service": service,
            "scope": scope
        }
        response = requests.get(auth_url, headers=headers, params=params)
    if response.status_code == 200:
        return response.json().get('token')
    return None

def get_manifest(registry_url, image_name, tag, token):
    headers = {
        "Authorization": f"Bearer {token}" if token else None
    }
    url = f"https://{registry_url}/v2/{image_name}/manifests/{tag}"
    response = requests.get(url, headers=headers)
    if response.status_code == 200:
        return response.json()
    return None

def download_blob(registry_url, image_name, digest, token):
    headers = {
        "Authorization": f"Bearer {token}" if token else None
    }
    url = f"https://{registry_url}/v2/{image_name}/blobs/{digest}"
    response = requests.get(url, headers=headers, stream=True)
    if response.status_code == 200:
        total_size = int(response.headers.get('content-length', 0))
        progress_bar = tqdm(total=total_size, unit='B', unit_scale=True, desc=f"Downloading {digest[:12]}")
        blob = b""
        for chunk in response.iter_content(chunk_size=1048576):  # 1MB chunk size
            if chunk:
                blob += chunk
                progress_bar.update(len(chunk))
        progress_bar.close()
        return blob
    return None

def download_layers(registry_url, image_name, layers, token):
    with ThreadPoolExecutor(max_workers=10) as executor:  # 多线程下载
        futures = [
            executor.submit(download_blob, registry_url, image_name, layer['digest'], token)
            for layer in layers
        ]
        return [future.result() for future in futures]

def save_blob_to_tar(tar, blob, blob_name):
    tarinfo = tarfile.TarInfo(name=blob_name)
    tarinfo.size = len(blob)
    tar.addfile(tarinfo, fileobj=io.BytesIO(blob))

def main():
    if len(sys.argv) < 3:
        print("Usage: python puller.py <image_url> <architecture> [<username> <password>]")
        sys.exit(1)

    image_url = sys.argv[1]
    architecture = sys.argv[2]
    username = sys.argv[3] if len(sys.argv) > 3 else None
    password = sys.argv[4] if len(sys.argv) > 4 else None

    # 解析镜像 URL
    if "/" not in image_url.split(":")[0]:
        # 如果没有指定仓库地址，默认使用 Docker Hub
        registry_url = DOCKER_HUB_REGISTRY_URL
        image_name = f"library/{image_url.split(':')[0]}"  # Docker Hub 的官方镜像需要加 library/
        tag = image_url.split(":")[1] if ":" in image_url else "latest"
    else:
        # 如果指定了仓库地址，使用指定的仓库
        if not image_url.startswith("http://") and not image_url.startswith("https://"):
            image_url = f"https://{image_url}"
        parsed_url = urlparse(image_url)
        registry_url = parsed_url.netloc
        image_name = parsed_url.path.lstrip('/').split(':')[0]
        tag = parsed_url.path.split(':')[1] if ':' in parsed_url.path else 'latest'

    # 获取认证 URL 和 Token
    auth_url = get_auth_url(registry_url)
    if not auth_url and registry_url != DOCKER_HUB_REGISTRY_URL:
        print("Failed to get auth URL")
        sys.exit(1)

    token = None
    if auth_url:
        token = get_auth_token(auth_url, username, password, registry_url, f"repository:{image_name}:pull")
        if not token and registry_url != DOCKER_HUB_REGISTRY_URL:
            print("Failed to get auth token")
            sys.exit(1)

    # 获取 Manifest
    manifest = get_manifest(registry_url, image_name, tag, token)
    if not manifest:
        print("Failed to get manifest")
        sys.exit(1)

    # 如果是多架构镜像，获取指定架构的 Manifest
    if manifest.get('mediaType') == 'application/vnd.docker.distribution.manifest.list.v2+json':
        for m in manifest.get('manifests', []):
            if m.get('platform', {}).get('architecture') == architecture:
                manifest = get_manifest(registry_url, image_name, m['digest'], token)
                break

    if not manifest:
        print("Failed to get architecture-specific manifest")
        sys.exit(1)

    # 打包镜像
    tar_filename = f"{image_name.replace('/', '_')}_{tag}_{architecture}.tar"
    with tarfile.open(tar_filename, 'w') as tar:
        # 生成 repositories 文件
        full_image_name = f"{registry_url}/{image_name}"  # 包含仓库地址
        repositories = {
            full_image_name: {
                tag: manifest['config']['digest']
            }
        }
        repositories_json = json.dumps(repositories, indent=2).encode()
        save_blob_to_tar(tar, repositories_json, 'repositories')

        # 生成 manifest.json 文件
        manifest_json = [
            {
                "Config": f"{manifest['config']['digest'].replace(':', '_')}.json",
                "RepoTags": [f"{full_image_name}:{tag}"],  # 包含仓库地址
                "Layers": [f"{layer['digest'].replace(':', '_')}.tar.gz" for layer in manifest['layers']]
            }
        ]
        manifest_json_bytes = json.dumps(manifest_json, indent=2).encode()
        save_blob_to_tar(tar, manifest_json_bytes, 'manifest.json')

        # 下载并保存各层文件
        layers = manifest.get('layers', [])
        blobs = download_layers(registry_url, image_name, layers, token)  # 多线程下载
        for blob, layer in zip(blobs, layers):
            if blob:
                save_blob_to_tar(tar, blob, layer['digest'].replace(':', '_') + '.tar.gz')

        # 下载并保存配置文件
        config_blob = download_blob(registry_url, image_name, manifest['config']['digest'], token)
        if config_blob:
            save_blob_to_tar(tar, config_blob, f"{manifest['config']['digest'].replace(':', '_')}.json")

    print(f"Image saved to {tar_filename}")

if __name__ == "__main__":
    main()